import pandas as pd
import plotly_express as px
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')
fig = px.choropleth(anx2010_df, locations="Code",
color="Anxiety_percent",
hover_name="Entity",
color_continuous_scale=px.colors.sequential.Plasma)
fig.show()Introduction
One of the important aspects of data analysis is making outputs that are easy to understand and aesthetically pleasing. In my work, I refer to this as making graphs “pretty”. Today I want to test whether LLMs have a concept of “prettiness”. I’m interested to see how they behave when I ask them to modify graph code to make it more aesthetically pleasing.
For reference, this is what the original graph looks like:
Test 1: GPT 4 Turbo
from openai import OpenAI
import os
import pandas as pd
from IPython.display import Markdown, display
from dotenv import load_dotenv
# Load API key
_ = load_dotenv()
client = OpenAI()
model = "gpt-4-turbo-2024-04-09"
system_role = '''You are a consultant skilled at making visually appealing and easy to understand
graphs to present to clients.'''
initial_prompt = '''Improve this code to make the graph better. Reply only with Python code.
Here is the code to improve:
import pandas as pd
import plotly_express as px
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')
fig = px.choropleth(anx2010_df, locations="Code",
color="Anxiety_percent",
hover_name="Entity",
color_continuous_scale=px.colors.sequential.Plasma)
fig.show()
'''
messages = [
{"role": "system", "content": system_role} ,
{"role": "user", "content": initial_prompt},
]
completion = client.chat.completions.create(
model=model, # the model to use
messages=messages, # a list of dictionaries with the prompt
n=3, # how many responses to generate
temperature=0.7 # temperature parameter
)
response = display(Markdown(f'-----\n#### Response {0} \n {resp.message.content} \n'))Response 0
import pandas as pd
import plotly_express as px
# Loading data
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')
# Creating a choropleth map with improved features
fig = px.choropleth(anx2010_df, locations="Code",
color="Anxiety_percent",
hover_name="Entity",
title="Global Anxiety Rates in 2010",
labels={"Anxiety_percent": "Anxiety Rate (%)"},
color_continuous_scale=px.colors.sequential.Plasma,
projection="natural earth")
# Enhancing layout
fig.update_layout(
geo=dict(
showframe=False,
showcoastlines=True,
projection_type='equirectangular'
),
coloraxis_colorbar=dict(
title="Anxiety Rate %",
thicknessmode="pixels", thickness=15,
lenmode="pixels", len=300
)
)
# Showing the plot
fig.show()Testing the code
import pandas as pd
import plotly_express as px
# Loading data
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')
# Creating a choropleth map with improved features
fig = px.choropleth(anx2010_df, locations="Code",
color="Anxiety_percent",
hover_name="Entity",
title="Global Anxiety Rates in 2010",
labels={"Anxiety_percent": "Anxiety Rate (%)"},
color_continuous_scale=px.colors.sequential.Plasma,
projection="natural earth")
# Enhancing layout
fig.update_layout(
geo=dict(
showframe=False,
showcoastlines=True,
projection_type='equirectangular'
),
coloraxis_colorbar=dict(
title="Anxiety Rate %",
thicknessmode="pixels", thickness=15,
lenmode="pixels", len=300
)
)
# Showing the plot
fig.show()Analysis
Here we can see that we have:
- A title for the overall graph
- A label for the anxiety rate legend
- No more black box around the outside of the map
Test 2: Claude 3.5 Sonnet
from dotenv import load_dotenv
from anthropic import Anthropic
import base64
import httpx
from IPython.display import display, Image, Markdown
# load api keys
_ = load_dotenv()
client = Anthropic()
message = client.messages.create(
model='claude-3-5-sonnet-20241022',
max_tokens=1000,
temperature=0.0,
system='''You are a consultant skilled at making visually appealing and easy to understand
graphs to present to clients.''',
messages=[
{"role": "user", "content": '''Improve this code to make the graph better. Reply only with Python code.
Here is the code to improve:
import pandas as pd
import plotly_express as px
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')
fig = px.choropleth(anx2010_df, locations="Code",
color="Anxiety_percent",
hover_name="Entity",
color_continuous_scale=px.colors.sequential.Plasma)
fig.show()
'''}
]
)
response = display(Markdown(f'-----\n#### Response {0} \n {message.content[0].text} \n'))Response 0
import pandas as pd
import plotly_express as px
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')
fig = px.choropleth(anx2010_df,
locations="Code",
color="Anxiety_percent",
hover_name="Entity",
hover_data={"Anxiety_percent": ":.2f"},
color_continuous_scale="RdYlBu_r",
title="Global Anxiety Prevalence (2010)",
labels={"Anxiety_percent": "Anxiety Prevalence (%)",
"Code": "Country Code",
"Entity": "Country"})
fig.update_layout(
title_x=0.5,
title_font_size=24,
geo=dict(
showframe=False,
showcoastlines=True,
projection_type='equirectangular'
),
width=1200,
height=700,
margin=dict(t=50, l=0, r=0, b=0)
)
fig.update_coloraxes(
colorbar_title_text="Prevalence (%)",
colorbar_title_side="right",
colorbar_thickness=15,
)
fig.show()Testing the code
import pandas as pd
import plotly_express as px
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')
fig = px.choropleth(anx2010_df,
locations="Code",
color="Anxiety_percent",
hover_name="Entity",
hover_data={"Anxiety_percent": ":.2f"},
color_continuous_scale="RdYlBu_r",
title="Global Anxiety Prevalence (2010)",
labels={"Anxiety_percent": "Anxiety Prevalence (%)",
"Code": "Country Code",
"Entity": "Country"})
fig.update_layout(
title_x=0.5,
title_font_size=24,
geo=dict(
showframe=False,
showcoastlines=True,
projection_type='equirectangular'
),
width=1200,
height=700,
margin=dict(t=50, l=0, r=0, b=0)
)
fig.update_coloraxes(
colorbar_title_text="Prevalence (%)",
colorbar_title_side="right",
colorbar_thickness=15,
)
fig.show() Analysis
Here we see:
- A new color scheme
- Updated titles
Conclusion
Both models did pretty well with this task. A lot of a graph’s “prettiness” comes down to personal preference, but asking LLMs for various options could definitely help in the brainstorming process.